package org.example.data_stream;

import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

import java.util.ArrayList;
import java.util.List;

/**
 * @author shenguangyang
 */
public class E10_TransformPartition {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        List<Event> events = new ArrayList<Event>() {{
            add(new Event("user02", "./cart", 2000L));
            add(new Event("user02", "./prod?id=10`", 3800L));
            add(new Event("user01", "./home1", 1000L));
            add(new Event("user02", "./prod?id=2`", 3000L));
            add(new Event("user01", "./prod?id=15`", 4000L));
            add(new Event("user03", "./prod?id=20`", 4000L));
            add(new Event("user04", "./prod?id=22`", 4000L));
            add(new Event("user05", "./prod?id=24`", 4000L));
        }};
        DataStream<Event> streamSource = env.fromCollection(events);

        // 1. 随机分区: shuffle 均匀分布到下游并行任务中
//        streamSource.shuffle().print().setParallelism(4);

        // 2. 轮训分区
//        streamSource.rebalance().print().setParallelism(4);
//        streamSource.print().setParallelism(4);

        // 3. rescale重缩放分区
//        DataStreamSource<Integer> streamSource1 = env.addSource(new MyRichParallelSourceFunction()).setParallelism(2);
//        streamSource1.rescale().print().setParallelism(4);

        // 4. 广播
//        streamSource.broadcast().print().setParallelism(4);

        // 5. 全局分区, 把所有数据都分到一个分区中
//        streamSource.global().print().setParallelism(4);

        // 6. 自定义重分区
        DataStream<Integer> stream2 = env.fromElements(1, 2, 3, 4, 5, 6, 7, 8)
                .partitionCustom(new MyPartitioner(), (KeySelector<Integer, Integer>) v -> v);
        stream2.print().setParallelism(4);

        env.execute();
    }

    private static class MyRichParallelSourceFunction extends RichParallelSourceFunction<Integer> {
        @Override
        public void run(SourceContext<Integer> ctx) throws Exception {
            for (int i = 0; i < 8; i++) {
                // 将奇数和偶数分别发送到0号和1号并行分区
                if (getRuntimeContext().getIndexOfThisSubtask() == i % 2) {
                    ctx.collect(i);
                }
            }
        }

        @Override
        public void cancel() {

        }
    }

    private static class MyPartitioner implements Partitioner<Integer> {
        @Override
        public int partition(Integer v, int i) {
            return v % 2;
        }
    }
}
